# -*- coding: utf-8 -*-
import asyncio
from dataclasses import dataclass
from typing import List, Any, Optional

import cv2
import mediapipe as mp

# MediaPipe Tasks - Pose Landmarker (多人支持)
from mediapipe.tasks.python.vision import PoseLandmarker, PoseLandmarkerOptions, RunningMode
from mediapipe.tasks.python.core import base_options as mp_base

# 平滑滤波
# from vision.smoother import LandmarkSmoother
from vision.smoother import AdaptiveLandmarkSmoother

# mp.Image 用于 Tasks 输入
mp_image = mp.Image


@dataclass
class FrameResult:
    """
    persons: List of persons, each person is a list of Pose landmarks
             结构对齐 Manager 期望：list[list[landmark-like]]
    frame_bgr: 当前帧 BGR 图像（OpenCV 格式）
    """
    persons: List[Any]
    frame_bgr: Any


class MediaPipeRunner:
    """
    多人版 Runner：使用 MediaPipe Tasks PoseLandmarker（VIDEO 模式）
    - 仅负责输出原始 landmarks 与帧
    - 稳定 person_id 的分配与维护由 Manager 统一负责
    """

    def __init__(
        self,
        camera_index: int = 0,
        model_path: str = "models/pose_landmarker_full.task",
        max_persons: int = 4,
        min_det_conf: float = 0.5,
        min_pres_conf: float = 0.5,
        min_track_conf: float = 0.5,
        width: int = 640,
        height: int = 360,
        fps: int = 30
    ):
        self.camera_index = camera_index
        self.model_path = model_path
        self.max_persons = max(1, min(8, int(max_persons)))  # 安全限制 1~8
        self.min_det_conf = float(min_det_conf)
        self.min_pres_conf = float(min_pres_conf)
        self.min_track_conf = float(min_track_conf)

        self.width = int(width)
        self.height = int(height)
        self.fps = int(fps)

        self.cap: Optional[cv2.VideoCapture] = None
        self.landmarker: Optional[PoseLandmarker] = None
        self._ts_ms: int = 0  # VIDEO 模式需要递增的 timestamp_ms

        # self.smoother = LandmarkSmoother(alpha = alpha)
        self.smoother = AdaptiveLandmarkSmoother(alpha_min=0.1, alpha_max=0.8, sensitivity=4.0, freeze_threshold=0.008, freeze_patience=3)

    # ---------- 生命周期 ----------
    def _open(self):
        # 摄像头
        self.cap = cv2.VideoCapture(self.camera_index)
        # 基础参数（可根据设备情况调整）
        if self.width > 0:
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        if self.height > 0:
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        if self.fps > 0:
            self.cap.set(cv2.CAP_PROP_FPS, self.fps)

        # Landmarker（VIDEO 模式）
        base = mp_base.BaseOptions(model_asset_path=self.model_path)
        options = PoseLandmarkerOptions(
            base_options=base,
            running_mode=RunningMode.VIDEO,  # 简单稳定，配合 timestamp_ms
            num_poses=self.max_persons,
            min_pose_detection_confidence=self.min_det_conf,
            min_pose_presence_confidence=self.min_pres_conf,
            min_tracking_confidence=self.min_track_conf,
        )
        self.landmarker = PoseLandmarker.create_from_options(options)
        self._ts_ms = 0

    def _close(self):
        if self.landmarker is not None:
            self.landmarker.close()
            self.landmarker = None
        if self.cap is not None:
            self.cap.release()
            self.cap = None

    # ---------- 推理 ----------
    def _detect_persons(self, frame_bgr) -> List[Any]:
        """
        将 OpenCV BGR 帧转换为 MediaPipe Image，并调用 landmarker
        返回：多人 landmarks 列表（每人是 list[landmark]）
        """
        # BGR -> RGB
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)

        # 构造 mp.Image（SRGB 格式）
        image = mp_image(image_format=mp.ImageFormat.SRGB, data=frame_rgb)

        # timestamp 递增（假定 fps ≈ self.fps）
        # 如果你有真实的时间间隔，建议用 time.monotonic() 计算更精确的间隔
        if self.fps > 0:
            self._ts_ms += int(1000 / self.fps)
        else:
            self._ts_ms += 33  # 兜底 30FPS

        # 执行检测
        result = self.landmarker.detect_for_video(image, self._ts_ms)

        persons: List[Any] = []
        if result is not None and getattr(result, "pose_landmarks", None):
            # result.pose_landmarks: List[List[landmark]]
            for lm_list in result.pose_landmarks:
                # 直接追加 landmarks 列表（NormalizedLandmark-like）
                persons.append(lm_list)

        persons = self.smoother.smooth(persons)

        return persons

    # ---------- 异步帧生成 ----------
    async def run(self):
        """
        异步帧生成器：
        - 打开资源
        - 循环读取帧、检测、产出 FrameResult
        - 调用方可：for await in runner.run(): manager.update_latest_poses(fr.persons)
        """
        self._open()
        try:
            while True:
                ok, frame = self.cap.read()
                if not ok:
                    # 读帧失败，稍作等待后重试
                    await asyncio.sleep(0.01)
                    continue

                persons = self._detect_persons(frame)
                yield FrameResult(persons=persons, frame_bgr=frame)

                # 将控制权交回事件循环，避免阻塞
                await asyncio.sleep(0)
        finally:
            self._close()